'''Genotype with Quintara Sequencing output .ab1 files
	Conventions: Samples are S#.N, where # is a construct, N is a clone
	Inputs: Construct Map, directory of sequencing files, (optional) thresholds for good reads.
	Output: Directory of high signal reads for simultaneous alignment in ApE
	WishList: Output: MUSCLE alignments of reads, +/- per clone, explanation of issues.'''

from Bio import Seq,SeqRecord,SeqIO,Align,AlignIO
from sequencing_auto import move_downloads
from numpy import array, max
import re
import csv
import os
import pdb
import sys
import subprocess

def parseKey(key):
    '''Return dictionary from CSV input. Invert primer dictionary.'''
    d, d_anti = {}, {}
    for row in key:
        if not row[0].count('>'):
            d[row[0]]=row[1]
    #Invert primer dictionaries (can index on primer number)
    d_inv = d.copy()
    for i,type in enumerate(['UTR','ORF','Anti']):
        for val in d[type].split(' '):
            d_inv[val]=i+1
    for i,type in enumerate(['Nat','Kan','Hyg']):
        for val in d[type].split(' '):
            d_anti[val]=type
    return d_inv, d_anti

def parseName(name):
    '''Map name to sample: [id, index, primer] '''
    pattern = '(S\d+)[.-](\d+)_(\w+)_'
    match = re.search(pattern,name)
    return list(match.groups())

def parseBlast(blast):
    '''Get most represented hit.'''
    headers = ''
    for line in blast.readlines():
        if line.count('>'):
            headers = headers + line
    maxgene = ''
    maxcounts = 0
    for gene in ['Bar1','Far1','Aga1','Aga2','Ace2','Nat','Kan','Hyg']:
        if headers.upper().count(gene.upper()) > maxcounts:
            maxgene = gene
    return maxgene
        

def genotype(seq_dir,d_anti,phred,minLength):
    '''Move through ab1 files, blastn, map to dictionary, write genotypes.
        Appends result to the index given by the sample name string. [id, index, primer, result]
        Sample: Locus 1 (UTR ORF ANTI) Locus 2 (UTR ORF ANTI) ...'''
    g = []
    blastOut=''
    files = os.listdir(seq_dir)
    #Keep only .ab1 files
    files = [ filename for i, filename in enumerate(files) if filename.endswith('.ab1') ]
    blasted = 0
    for i, filename in enumerate(files): 
        #Check for current sample_id (case-insensitive)
        f = open(seq_dir+filename,'rb')
        seq_iter = SeqIO.AbiIO.AbiIterator(f,trim=phred)
        #Add Mott-trimmed sequence, P(error) < phred (Default = 0.01)
        seq = seq_iter.next()
        if len(seq) > minLength:
            #Write trimmed fasta
            SeqIO.write(seq,seq.name+'.fasta','fasta')
            #Parse name for data table index
            index = parseName(seq.name)
            #If antibiotic primer, write antibiotic, else Blast
            if d_anti.has_key(index[2]):
                index.append(d_anti[index[2]])
            else:
                blastOut = open('blast.out','w')
                subprocess.call(['blastall','-p','blastn','-i',seq.name+'.fasta','-d','/data/star/blastdb/nt','-e','0.00001','-a','32','-W','20','-v','5'],stdout=blastOut)
                blastOut = open('blast.out','r')
                blastHit = parseBlast(blastOut)
                index = parseName(seq.name)
                index.append(blastHit)
            g.append(index)
            print str(i/float(len(files))*100)+'% Complete.'
            blasted = blasted+1
        f.close()
    print g
    print 'Examined '+str(len(files))+' files.'
    print 'Blasted '+str(blasted)+' sequences.'
    return g

def writeResults(d,g):
    '''Output as table. Input is [S#,Index,Primer,data]. Use dictionary for table row on strain #.'''
    mapped, numonly = [], []
    for row in g:
        mapped.append([int(row[0][1:]),int(row[1]),int(d[row[2]]+1),row[3]])
        numonly.append([int(row[0][1:]),int(row[1]),int(d[row[2]]+1)])
    g_array = array(mapped)
    n_array = array(numonly)
    #Initialize
    out = [ ['']*max(n_array[:,1])*max(n_array[:,2]) for x in range(max(n_array[:,0])) ]
    for row in g_array:
        out[int(row[0])-1][(int(row[1])-1)*max(n_array[:,2])+int(row[2])-1] = str(out[int(row[0])-1][(int(row[1])-1)*max(n_array[:,2])+int(row[2])-1])+' '+row[3]
    fout = open(seq_dir+'genotyping_results.csv','w')
    fout.write('Sample,Strain yTS-,Bar1,Bar1,Bar1,,Far1,Far1,Far1,,Aga1/Ace2,Aga1/Ace2,Aga1/Ace2,,Aga2,Aga2,Aga2\n')
    fout.write(',,UTR,ORF,Anti,,UTR,ORF,Anti,,UTR,ORF,Anti,,UTR,ORF,Anti\n')
    for i, row in enumerate(out):
        fout.write(str(i+1)+','+d['S'+str(i+1)]+','.join(row)+'\n')
        print '\t'.join(row)
    fout.close()
    print 'Wrote'+fout.name 
    return out
    

def main(argv):
    '''USAGE: python genotype_auto.py [sequencing_directory] [phred_cutoff] [minimum_length_of_read]'''
    if len(argv)<4:
        minLength = 500
        if len(argv)<3:
            phred = 0.01
            if len(argv)<2:    
                seq_dir = '../Dropbox/Yeast_Clustering/Sequencing/Sequencing_111112/'
            else:   
                seq_dir = argv[1]
        else:
            seq_dir = argv[1]
            phred = argv[2]
            phred = float(phred)
    else:
        script, seq_dir, phred, minLength = argv
        phred = float(phred)
        minLength = int(minLength)

    seq_dir = './'
    #zips = move_downloads(seq_dir)
    key = csv.reader(open(seq_dir+'key.csv','rU'), delimiter=',')
    dictionary, dict_anti = parseKey(key)
    genotyped = genotype(seq_dir,dict_anti,phred,minLength)
    out = writeResults(dictionary,genotyped)

if __name__ == "__main__":
    sys.exit(main(sys.argv))